#!/usr/bin/env python
#
# frame_extractor.py takes tgrep2 s-expression output
# and outputs verb sub-categorization frames.
#
# Suggested usage:
#
# First extract matched subtrees as s-expressions:
# ~> tgrep2 -i -a -c wsj2.t2c 'VP < (__</^abandon[esi]|^abandon$/)' >abandon.par
#
# Next run:
#
# ~> frame_extractor.py -p abandon.par
#
# Example:
# ./frame_extractor.py -p abandon.par |sort|uniq -c|sort -nr
# will output:
# 12 [VP (VB abandon) NP]
#  9 [VP (VBN abandoned) NP]
#  7 [VP (VBD abandoned) NP]
#  6 [VP (VBG abandoning) NP]
#  2 [VP (VBZ abandons) NP]
#  2 [VP (VBN abandoned) NP PP PP-TMP]
#  1 [VP (VBP abandon) NP]
#  1 [VP (VBN abandoned) NP SBAR-TMP]
# ...
#
# ./frame_extractor.py -p abandon.par -x true
# will output the frame with the arguments expanded to words:
#
# [VP (VB abandon) (NP the provision initially drafted * by Mr. Gray)]
# [VP (VBN abandoned) (NP *) (PP by the company) (PP-TMP in the 1980s)]
# ...
#
# author: yks
#
import re
import sys
import nltk
from optparse import OptionParser

def find_verb(tree):
    for i in xrange(len(tree)):
        if re.match("^VB.*$", tree[i].node):
            return tree[i],i
    return [None, None]
        
def main():
    parser = OptionParser()
    parser.add_option("-p", "--parses", dest="parses",
                      help="a file of parse trees s-expression", metavar="file")

    parser.add_option("-x", "--expand", dest="expand",
                      help="expand the argument of the verb frame")

    (options, args) = parser.parse_args()

    if options.parses is None:
        parser.print_help()
        exit(1)
        
    expand = False
    if options.expand is not None:
        expand = options.expand
        
    parses = []
    f = open(options.parses,"r")
    parses = f.readlines()
    parses = [parse.strip() for parse in parses]
    f.close()

    for parse in parses:
        tree = nltk.bracket_parse(parse)
        verb,index = find_verb(tree)
        if not verb:
            continue
        tag = verb.node
        prefix = tree[0:index]
        word = verb[0].lower()
        args = tree[index+1:]
        if expand:
            rest = ["(%s %s)" %(arg.node, " ".join(arg.leaves())) for arg in args]
        else:
            rest = [arg.node for arg in args]

        s = "[%s (%s %s) %s]" %(tree.node, tag, word, " ".join(rest))
        print s

if __name__ == '__main__': main()
